import requests
import concurrent
from concurrent.futures import ThreadPoolExecutor
from re import search as re_search
from yaml import safe_load as load_yaml
from urllib3 import disable_warnings
from commons.encrypt import encrypt_password
from datetime import datetime, timedelta
from config import settings

THREAD_LIMIT = 15
UA = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/115.0.0.0 Safari/537.36"
headers_init = {
    "User-Agent": UA,
    "Upgrade-Insecure-Requests": "1"
}


def headers_with_referer(url):
    return {
        "User-Agent": UA,
        "Upgrade-Insecure-Requests": "1",
        "Referer": url
    }


def fmt(time_obj):
    return time_obj.strftime("%m.%d %H:%M:%S")


# 登录函数
def login_(url, username, password):
    disable_warnings()
    session = requests.session()
    init_page = session.get(url, headers=headers_init, verify=False)

    set_cookie = init_page.headers.get('Set-Cookie', None)
    assert set_cookie, 'Set-Cookie为空'
    CSRF = re_search(r"login_csrftoken=(.*?);", set_cookie)
    assert CSRF, 'CSRF-Token不存在'
    RSA_PUBLIC_KEY = re_search(r'C_KEY = "(.*?)"', init_page.text)
    assert RSA_PUBLIC_KEY, 'RSA公钥不存在'
    login_data = {
        "username": username,
        "password": encrypt_password(password, RSA_PUBLIC_KEY.group(1)),
        "csrfmiddlewaretoken": CSRF.group(1),
    }
    headers = headers_with_referer(url)

    login_resp = session.post(url, data=login_data, headers=headers, verify=False)
    assert login_resp.status_code < 400, f'状态码：{login_resp.status_code}'
    assert 'ang="en">' not in login_resp.text[0:60], '用户名或密码错误'

    redirect_resp = session.get(login_resp.url, headers=headers, verify=False)
    assert redirect_resp.status_code < 400, f'重定向失败：{login_resp.url}'

    cookies = session.cookies.get_dict()
    assert cookies, 'Cookie为空'

    return ';'.join([f'{k}={v}' for k, v in cookies.items()])


# 批量登录并保存cookie
def concurrent_login(url, user_yaml, output_file):
    users = {}
    # user_yaml 支持传入字符串或字符串列表
    bound = user_yaml if isinstance(user_yaml, list) else [user_yaml]
    for file in bound:
        with open(file, 'r') as f:
            users.update(load_yaml(f))
    data = {}
    print('————————登录开始————————')
    begin = datetime.now()
    # 使用线程池来并发执行登录操作
    with ThreadPoolExecutor(THREAD_LIMIT) as TPE:
        futures = {TPE.submit(login_, url, user, pwd): user for user, pwd in users.items()}
        for future in concurrent.futures.as_completed(futures):
            user = futures[future]
            try:
                cookie = future.result()  # 获取线程执行的结果
                data[user] = cookie
            except Exception as e:
                print(f"登录失败用户：{user} {str(e)}")
    end = datetime.now()
    print('————————登录结束————————')
    print(f'登录时间：{fmt(begin)}')
    print(f'完成时间：{fmt(end)}')
    print(f'失效时间：{fmt(begin + timedelta(hours=3))}')
    # 登录操作完成后，将结果写入文件
    with open(output_file, 'w') as file:
        for user, cookie in data.items():
            file.write(f"{user},{cookie}\n")


if __name__ == "__main__":
    cfg = settings.cfg
    PARAMS = {
        "url": cfg.get('itsm').get('itsm_login_url'),
        "output_file": "D:/项目资料/ITSM/国能-性能测试/user.txt",
        "user_yaml": [
            "./user_files/search_users.yaml",
        ],
    }
    concurrent_login(**PARAMS)
